package main

// hash and math libraries
import (
	"crypto/sha256"
	"fmt"
	"github.com/consensys/gnark-crypto/ecc"
	bn254 "github.com/consensys/gnark-crypto/ecc/bn254/fr/mimc"
	"github.com/consensys/gnark/backend/groth16"
	"github.com/consensys/gnark/frontend"
	r1cs2 "github.com/consensys/gnark/frontend/cs/r1cs"
	"github.com/xuri/excelize/v2"
	"hash"
	"io"
	"log"
	"os"
	"strconv"
	"time"
)

import (
	"math/big"
)

// pairing-based cryptography library
import (
	"github.com/Nik-U/pbc"
)

type Kzg_scheme struct {
	// The number of blocks generated by the file
	chunckNum int
	// The number of chunks in each block
	blockNum int
	// The tags for the file blocks
	pairing *pbc.Pairing
	// The public parameters for the KZG scheme
	_g      *pbc.Element
	_h      *pbc.Element
	_xi     *pbc.Element
	epsilon *pbc.Element
	_Dx     [][][]byte
	hashDx  [][]*pbc.Element
	hash    hash.Hash
	phi     []*pbc.Element
	eta     []*pbc.Element
	_theta  []*pbc.Element
	_acc    *pbc.Element
}

func NewKzgScheme(Mx [][][]byte) *Kzg_scheme {
	row := len(Mx)
	collumn := len(Mx[0])
	chunckNum := row
	blockNum := collumn
	_g := pairing.NewG1().Rand()
	_h := pairing.NewG2().Rand()
	kzg_hash := sha256.New()
	epsilon := pairing.NewZr().Rand()
	_xi := pairing.NewZr().Rand()
	phi := make([]*pbc.Element, collumn+1)
	for i := 0; i < collumn+1; i++ {
		newi := pairing.NewZr().SetInt32(int32(i))
		phi_i := pairing.NewG1().PowZn(_g, pairing.NewZr().PowZn(_xi, newi))
		phi[i] = phi_i
	}
	eta := make([]*pbc.Element, row+1)
	for i := 0; i < row+1; i++ {
		newi := pairing.NewZr().SetInt32(int32(i))
		epsilon_i := pairing.NewG1().PowZn(_h, pairing.NewZr().PowZn(epsilon, newi))
		eta[i] = epsilon_i
	}
	_theta := make([]*pbc.Element, chunckNum)
	for i := 0; i < chunckNum; i++ {
		_theta[i] = pairing.NewZr().Set0()
	}

	return &Kzg_scheme{
		chunckNum: chunckNum,
		blockNum:  blockNum,
		pairing:   pairing,
		_g:        _g,
		_h:        _h,
		_xi:       _xi,
		epsilon:   epsilon,
		_Dx:       Mx,
		hash:      kzg_hash,
		phi:       phi,
		eta:       eta,
		_theta:    _theta,
	}
}

func mimcHash(data []byte) string {
	f := bn254.NewMiMC()
	f.Write(data)
	//fmt.Println(n)
	//fmt.Println(err)
	hash := f.Sum(nil)
	hashInt := big.NewInt(0).SetBytes(hash)
	return hashInt.String()
}

func (kzg *Kzg_scheme) tag_Gen() {
	hashMx := kzg._Dx
	HMx := make([][]*pbc.Element, kzg.chunckNum)
	//new_x := big.NewInt(1)
	//User_HMx_all := pairing.NewZr().Set0()
	t_tagHash = time.Now()
	for i := 0; i < len(hashMx); i++ {
		HMx[i] = make([]*pbc.Element, kzg.blockNum)
		for j := 0; j < len(hashMx[0]); j++ {
			t := mimcHash(hashMx[i][j])
			//t := hashMx[i][j]
			HMx[i][j], _ = pairing.NewZr().SetString(t, 10)
			//v, _ := pairing.NewZr().SetString(t, 10)
			//User_HMx_all.Add(User_HMx_all, v)
		}
	}
	//fmt.Println(User_HMx_all)
	t_tagHash_end = time.Now()
	//log.Println("tagGen HMx time:", t_tagHash_end.Sub(t_tagHash))

	//signatures := make([]*pbc.Element, kzg.chunckNum)
	for i := 0; i < len(hashMx); i++ {
		//t_eachTag := time.Now()
		Lixi := pbcPolynomialiEval(HMx, kzg._xi, i)
		g_Lixi := pairing.NewG1().PowZn(kzg._g, Lixi)
		//t_eachTag_end := time.Now()
		//log.Println(t_eachTag_end.Sub(t_eachTag))
		kzg._theta[i] = g_Lixi
	}
	prod := pairing.NewZr().Set1()
	for i := 0; i < len(kzg._theta); i++ {
		ZHtheta := pairing.NewZr().SetFromStringHash(kzg._theta[i].String(), kzg.hash)
		prod = pairing.NewZr().Mul(prod, pairing.NewZr().Add(ZHtheta, kzg.epsilon))
	}
	kzg._acc = pairing.NewG2().PowZn(kzg._h, prod)

}

func (kzg *Kzg_scheme) challenge(CNum int) map[int]*pbc.Element {
	if CNum < 0 || CNum > kzg.chunckNum+1 {
		panic("CNum wrong!")
	}
	selectBlockSet := getUniqueRandomNum(kzg.chunckNum, CNum)
	chal := make(map[int]*pbc.Element)
	comm_chal := 0
	for _, i := range selectBlockSet {
		chal[i] = pairing.NewZr().Rand()
		comm_chal += chal[i].BytesLen()
	}
	//log.Println("comm_chal:", comm_chal)
	return chal
}

func (kzg *Kzg_scheme) hashMxGen(chal map[int]*pbc.Element) ([][]*big.Int, *big.Int) {
	hashMx := kzg._Dx
	HMx := make([][]*big.Int, kzg.chunckNum)
	HMx_all := big.NewInt(0)
	temp := big.NewInt(0)
	//new_x := big.NewInt(1)
	for i := range chal {
		HMx[i] = make([]*big.Int, kzg.blockNum)
		for j := 0; j < len(hashMx[i]); j++ {
			t := mimcHash(hashMx[i][j])
			//t := hashMx[i][j]
			//newt := pairing.NewZr().SetBytes(t)
			s, _ := big.NewInt(0).SetString(t, 10)
			juhe := big.NewInt(0).Mul(chal[i].BigInt(), s)
			//newt, _ := pairing.NewZr().SetString(t, 10)
			//juhe := pairing.NewZr().Mul(chal[i], newt)
			//juhe = juhe.Mul(v,r.BigInt())
			HMx[i][j] = juhe
			HMx_all.Add(HMx_all, juhe)
			temp.Add(temp, juhe)
		}
	}
	return HMx, HMx_all
}

func (kzg *Kzg_scheme) proofGen(chal map[int]*pbc.Element, HMx [][]*big.Int) (map[int]*pbc.Element, *pbc.Element, []*pbc.Element) {
	Lx_sum := make([]*pbc.Element, kzg.blockNum)
	for i := range Lx_sum {
		Lx_sum[i] = pairing.NewZr().Set0()
	}
	//gamma := pairing.NewZr().Rand()
	//Gamma := pairing.NewG1().PowZn(kzg._g, gamma)
	//alpha := pairing.NewZr().SetFromStringHash(Gamma.String(), kzg.hash)
	for i, _ := range chal {
		for j := 0; j < kzg.blockNum; j++ {
			newHMx := pairing.NewZr().SetBig(HMx[i][j])
			Lx_sum[j] = pairing.NewZr().Add(Lx_sum[j], newHMx)
		}
	}
	//Lx_sum[0].Add(Lx_sum[0], pairing.NewZr().Mul(alpha, gamma))
	//Lu := polynomialEval(Lx_sum, u)

	//Lu_slice := make([]*pbc.Element, 1)
	//Lu_slice[0] = Lu
	//Lx_Lu := make([]*pbc.Element, len(Lx_sum))
	//copy(Lx_Lu, Lx_sum)
	//Lx_Lu = polynomialSub(Lx_Lu, Lu_slice)
	//
	//x_u := make([]*pbc.Element, 2)
	//x_u[0] = fNeg(u)                     //-r
	//x_u[1] = pairing.NewZr().Set1()      //1
	//Qx, rem := polynomialDiv(Lx_Lu, x_u) //pkx-pkr/-r,1
	//for i := 0; i < len(rem); i++ {
	//	if !rem[i].Is0() {
	//		fmt.Println("rem is not zero")
	//	}
	//}
	wit := pairing.NewG2().Set1()
	theta_prf := make(map[int]*pbc.Element)
	dictD_C := make(map[int]*pbc.Element)

	for i := 0; i < kzg.chunckNum; i++ {
		dictD_C[i] = kzg._theta[i]
	}
	for i := range chal {
		delete(dictD_C, i)
		theta_prf[i] = kzg._theta[i]
	}

	prodD_C := make([]*pbc.Element, 1)
	prodD_C[0] = pairing.NewZr().Set1()
	for i := range dictD_C {
		slice := make([]*pbc.Element, 2)
		slice[0] = pairing.NewZr().SetFromStringHash(dictD_C[i].String(), kzg.hash)
		slice[1] = pairing.NewZr().Set1()
		prodD_C = polynomialMul(prodD_C, slice)
	}
	for i := 0; i < len(prodD_C); i++ {
		wit.Mul(wit, pairing.NewG1().PowZn(kzg.eta[i], prodD_C[i]))
	}

	return theta_prf, wit, Lx_sum
}
func (kzg *Kzg_scheme) verify(chal map[int]*pbc.Element, theta_prf map[int]*pbc.Element, wit *pbc.Element, Lx_sum []*pbc.Element) bool {
	//alpha := pairing.NewZr().SetFromStringHash(Gamma.String(), kzg.hash)
	left1 := pairing.NewGT().Pair(kzg._acc, kzg._h)
	prod := pairing.NewG2().Set1()
	//prod := pairing.NewZr().Set1()
	left2 := pairing.NewG1().Set1()
	//total_eta := pairing.NewG2().Set1()
	//t_left := time.Now()
	for i := range chal {
		//zthetai := pairing.NewZr().SetFromStringHash(kzg._theta[i].String(), kzg.hash)
		//prod.Mul(prod, pairing.NewZr().Add(zthetai, kzg.epsilon))
		left2.Mul(left2, pairing.NewG1().PowZn(theta_prf[i], chal[i]))
	}

	//-------------------------------------------------------------------
	// on-chain auditing
	dictC := make(map[int]*pbc.Element)

	for i := range chal {
		dictC[i] = kzg._theta[i]
	}

	prodC := make([]*pbc.Element, 1)
	prodC[0] = pairing.NewZr().Set1()
	for i := range dictC {
		slice := make([]*pbc.Element, 2)
		slice[0] = pairing.NewZr().SetFromStringHash(dictC[i].String(), kzg.hash)
		slice[1] = pairing.NewZr().Set1()
		prodC = polynomialMul(prodC, slice)
	}
	for i := 0; i < len(prodC); i++ {
		prod.Mul(prod, pairing.NewG1().PowZn(kzg.eta[i], prodC[i]))
	}
	//------------------------------------------------------------------
	e4 := prod
	//e4 := pairing.NewG2().PowZn(kzg._h, prod)
	right1 := pairing.NewGT().Pair(wit, e4)

	right2 := pairing.NewG1().Set1()
	for i := 0; i < kzg.blockNum; i++ {
		right2.Mul(right2, pairing.NewG1().PowZn(kzg.phi[i], Lx_sum[i]))
	}

	px := polynomialEval(Lx_sum, kzg._xi)
	g_px := pairing.NewG1().PowZn(kzg._g, px)
	g_px.Equals(left2)
	left1.Equals(right1)
	return left2.Equals(right2)
}

func init() {
	// 获取日志文件句柄
	// 以 只写入文件|没有时创建|文件尾部追加 的形式打开这个文件
	logFile, err := os.OpenFile(`./comm3.txt`, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		panic(err)
	}
	// 设置存储位置
	log.SetOutput(logFile)
	// 组合一下即可，os.Stdout代表标准输出流
	multiWriter := io.MultiWriter(os.Stdout, logFile)
	log.SetOutput(multiWriter)

	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
}

type excel struct {
	chunckNum      int           `json:"ChunckNum"`
	blockNum       int           `json:"BlockNum"`
	chalNum        int           `json:"ChalNum"`
	blockSize      int           `json:"BlockSize"`
	filePolyGen    time.Duration `json:"FilePolyGen"`
	Initialize     time.Duration `json:"Initialize"`
	tagGenHMx      time.Duration `json:"TagGenHMx"`
	TagGen         time.Duration `json:"TagGen"`
	challenge      time.Duration `json:"ChalTime"`
	hashMxGen      time.Duration `json:"HashMxGen"`
	zksnarkCompile time.Duration `json:"SnarkCompile"`
	r1csSetup      time.Duration `json:"r1csSetup"`
	zkProof        time.Duration `json:"zkProof"`
	zkVerify       time.Duration `json:"zkVerify"`
	zkTime         time.Duration `json:"zkTime"`
	ProofGen       time.Duration `json:"ProofGen"`
	verify         time.Duration `json:"Verify"`
	allTime        time.Duration `json:"allTime"`
}

func main() {
	for g := 0; g < 1; g++ {
		log.Println("chunkNum", chunckNum)
		log.Println("blockNum", blockNum)
		t_start := time.Now()
		t_fileGen := time.Now()
		Mx := hashFilePolyGen2(chunckNum, 0, blockNum)
		t_fileGen_end := time.Now()
		log.Println("file_poly_gen time: ", t_fileGen_end.Sub(t_fileGen))

		t_init := time.Now()
		kzgSchemeTest := NewKzgScheme(Mx)
		t_init_end := time.Now()
		log.Println("Initialize time: ", t_init_end.Sub(t_init))

		t_tag := time.Now()
		kzgSchemeTest.tag_Gen()
		t_tag_end := time.Now()
		log.Println("TagGen time: ", t_tag_end.Sub(t_tag))

		t_chal := time.Now()
		chal := kzgSchemeTest.challenge(chalNumber)
		t_chal_end := time.Now()
		log.Println("challenge time: ", t_chal_end.Sub(t_chal))

		t_hashMx := time.Now()
		HMx, HMx_all := kzgSchemeTest.hashMxGen(chal)
		t_hashMx_end := time.Now()

		t_proof := time.Now()
		theta_prf, wit, Lx_sum := kzgSchemeTest.proofGen(chal, HMx)
		t_proof_end := time.Now()
		log.Println("ProofGen time: ", t_proof_end.Sub(t_proof))

		t_verify := time.Now()
		result := kzgSchemeTest.verify(chal, theta_prf, wit, Lx_sum)
		t_verify_end := time.Now()

		log.Println("verify result:", result)

		tzk := time.Now()
		hashMx := kzgSchemeTest._Dx
		var circuit Circuit
		circuit.Mx = make([][]frontend.Variable, len(chal))
		circuit.chal = make([]frontend.Variable, len(chal))

		for i := range chal {
			circuit.Mx[i] = make([]frontend.Variable, len(hashMx[i]))
			circuit.chal[i] = chal[i].Bytes()
			for j := 0; j < len(hashMx[i]); j++ {
				circuit.Mx[i][j] = frontend.Variable(hashMx[i][j])
			}
		}

		t_compile := time.Now()
		r1cs, err := frontend.Compile(ecc.BN254, r1cs2.NewBuilder, &circuit)
		t_compile_end := time.Now()
		if err != nil {
			fmt.Printf("Compile failed : %v\n", err)
			return
		}

		t_r1cs := time.Now()
		pk, vk, err := groth16.Setup(r1cs)
		t_r1cs_end := time.Now()
		if err != nil {
			fmt.Printf("Setup failed\n")
			return
		}

		var assignment Circuit
		assignment.Mx = make([][]frontend.Variable, len(chal))
		circuit.chal = make([]frontend.Variable, len(chal))
		for i := range chal {
			assignment.Mx[i] = make([]frontend.Variable, len(hashMx[i]))
			circuit.chal[i] = chal[i].Bytes()
			for j := 0; j < len(hashMx[i]); j++ {
				assignment.Mx[i][j] = frontend.Variable(hashMx[i][j])
			}
		}
		assignment.HMx_all = frontend.Variable(HMx_all.Bytes())

		witness, err := frontend.NewWitness(&assignment, ecc.BN254)

		if err != nil {
			fmt.Println("witness:", err)
		}

		t_zkProof := time.Now()
		proof, err := groth16.Prove(r1cs, pk, witness)
		t_zkProof_end := time.Now()
		if err != nil {
			fmt.Printf("Prove failed： %v\n", err)
			return
		}
		comm_prove := 0
		for i := range Lx_sum {
			comm_prove += Lx_sum[i].BytesLen()
		}
		for i := range theta_prf {
			comm_prove += theta_prf[i].BytesLen()
		}
		comm_prove += wit.BytesLen()
		proofPath := "proof.gob"
		file, _ := os.Create(proofPath)
		defer file.Close()
		_, err = proof.WriteTo(file)
		if err != nil {
			return
		}
		fileInfo, _ := file.Stat()
		proofSize := fileInfo.Size()
		comm_prove += int(proofSize)
		log.Println("comm_prove:", comm_prove)
		publicWitness, err := witness.Public()
		//fmt.Println(err)
		if err != nil {
			fmt.Println("public witness:", err)
		}

		t_zkVerify := time.Now()
		err = groth16.Verify(proof, vk, publicWitness)
		t_zkVerify_end := time.Now()
		if err != nil {
			fmt.Println("verify:", err)
		}
		tzk_end := time.Now()
		//log.Println("zk time: ", tzk_end.Sub(tzk))
		t_end := time.Now()

		excel1 := &excel{
			chunckNum:      chunckNum,
			blockNum:       blockNum,
			blockSize:      bitNum,
			chalNum:        chalNumber,
			filePolyGen:    t_fileGen_end.Sub(t_fileGen),
			Initialize:     t_init_end.Sub(t_init),
			tagGenHMx:      t_tagHash_end.Sub(t_tagHash),
			TagGen:         t_tag_end.Sub(t_tag),
			challenge:      t_chal_end.Sub(t_chal),
			hashMxGen:      t_hashMx_end.Sub(t_hashMx),
			zksnarkCompile: t_compile_end.Sub(t_compile),
			r1csSetup:      t_r1cs_end.Sub(t_r1cs),
			zkProof:        t_zkProof_end.Sub(t_zkProof),
			zkVerify:       t_zkVerify_end.Sub(t_zkVerify),
			zkTime:         tzk_end.Sub(tzk),
			ProofGen:       t_proof_end.Sub(t_proof),
			verify:         t_verify_end.Sub(t_verify),
			allTime:        t_end.Sub(t_start),
		}
		if write == 1 {
			xlsx, err := excelize.OpenFile("multiblock4.xlsx")
			//xlsx := excelize.NewFile()
			xlsx.NewSheet("result")
			if err != nil {
				fmt.Println("excel error:", err)
			}
			o := 1
			err = xlsx.SetCellStr("result", "A"+strconv.Itoa(o), "chunckNum")
			if err != nil {
				fmt.Println(err)
			}
			xlsx.SetCellValue("result", "B"+strconv.Itoa(o), "blockNum")
			xlsx.SetCellValue("result", "C"+strconv.Itoa(o), "blockSize")
			xlsx.SetCellValue("result", "D"+strconv.Itoa(o), "chalNum")
			xlsx.SetCellValue("result", "E"+strconv.Itoa(o), "filePolyGen")
			xlsx.SetCellValue("result", "F"+strconv.Itoa(o), "Initialize")
			xlsx.SetCellValue("result", "G"+strconv.Itoa(o), "tagGenHMx")
			xlsx.SetCellValue("result", "H"+strconv.Itoa(o), "TagGen")
			xlsx.SetCellValue("result", "I"+strconv.Itoa(o), "challenge")
			xlsx.SetCellValue("result", "J"+strconv.Itoa(o), "hashMxGen")
			xlsx.SetCellValue("result", "K"+strconv.Itoa(o), "zksnarkCompile")
			xlsx.SetCellValue("result", "L"+strconv.Itoa(o), "r1csSetup")
			xlsx.SetCellValue("result", "M"+strconv.Itoa(o), "zkProof")
			xlsx.SetCellValue("result", "N"+strconv.Itoa(o), "zkVerify")
			xlsx.SetCellValue("result", "O"+strconv.Itoa(o), "zkTime")
			xlsx.SetCellValue("result", "P"+strconv.Itoa(o), "ProofGen")
			xlsx.SetCellValue("result", "Q"+strconv.Itoa(o), "verify")
			xlsx.SetCellValue("result", "R"+strconv.Itoa(o), "allTime")
			xlsx.SetCellValue("result", "S"+strconv.Itoa(o), "Store")
			xlsx.SetCellValue("result", "T"+strconv.Itoa(o), "PrfGen")
			xlsx.SetCellValue("result", "U"+strconv.Itoa(o), "PrfVeri")

			//xlsx, err := excelize.OpenFile("multiblock2.xlsx")
			rows, err := xlsx.GetRows("result")
			if err != nil {
				fmt.Println(err)
			}
			i := len(rows) + 1
			xlsx.SetCellValue("result", "A"+strconv.Itoa(i), chunckNum)
			xlsx.SetCellValue("result", "B"+strconv.Itoa(i), excel1.blockNum)
			xlsx.SetCellValue("result", "C"+strconv.Itoa(i), excel1.blockSize)
			xlsx.SetCellValue("result", "D"+strconv.Itoa(i), excel1.chalNum)
			xlsx.SetCellValue("result", "E"+strconv.Itoa(i), excel1.filePolyGen.String())
			xlsx.SetCellValue("result", "F"+strconv.Itoa(i), excel1.Initialize.String())
			xlsx.SetCellValue("result", "G"+strconv.Itoa(i), excel1.tagGenHMx.String())
			xlsx.SetCellValue("result", "H"+strconv.Itoa(i), excel1.TagGen.String())
			xlsx.SetCellValue("result", "I"+strconv.Itoa(i), excel1.challenge.String())
			xlsx.SetCellValue("result", "J"+strconv.Itoa(i), excel1.hashMxGen.String())
			xlsx.SetCellValue("result", "K"+strconv.Itoa(i), excel1.zksnarkCompile.String())
			xlsx.SetCellValue("result", "L"+strconv.Itoa(i), excel1.r1csSetup.String())
			xlsx.SetCellValue("result", "M"+strconv.Itoa(i), excel1.zkProof.String())
			xlsx.SetCellValue("result", "N"+strconv.Itoa(i), excel1.zkVerify.String())
			xlsx.SetCellValue("result", "O"+strconv.Itoa(i), excel1.zkTime.String())
			xlsx.SetCellValue("result", "P"+strconv.Itoa(i), excel1.ProofGen.String())
			xlsx.SetCellValue("result", "Q"+strconv.Itoa(i), excel1.verify.String())
			xlsx.SetCellValue("result", "R"+strconv.Itoa(i), excel1.allTime.String())
			xlsx.SetCellValue("result", "S"+strconv.Itoa(i), excel1.TagGen.String())
			xlsx.SetCellValue("result", "T"+strconv.Itoa(i), (excel1.ProofGen + excel1.zkProof).String())
			xlsx.SetCellValue("result", "U"+strconv.Itoa(i), (excel1.verify + excel1.zkVerify).String())

			err = xlsx.SaveAs("multiblock4.xlsx")
			if err != nil {
				fmt.Println("excel error", err)
			}
		}
	}
}
